﻿using System.Collections;
using System.Collections.Generic;
using UnityEngine;

namespace RVO3D
{
    public class KdTree
    {
        internal class AgentTreeNode
        {
            internal int begin;
            internal int end;
            internal int left;
            internal int right;
            internal Vector3 maxCoord;  //最大坐标
            internal Vector3 minCoord;   // 最小坐标
        }

        internal Agent[] agents_ ;
        internal AgentTreeNode[] agentTree_;



        /// <summary>
        /// 构建一个代理KD树
        /// </summary>
        internal void BuildAgentTree()
        {
            //agents_ = RVOSimulator.Instance.agents_;

            //if (agents_ != null)
            //{
            //    agentTree_ = new AgnetTreeNode[2 * agents_.Count -1];

            //    for (int i = 0; i < agents_.Count; i++)
            //    {
            //        agentTree_[i] = new AgnetTreeNode();
            //    }

            //    BuildAgentTreeRecursive(0, agents_.Count, 0);
            //}

            if (agents_ == null || agents_.Length != RVOSimulator.Instance.agents_.Count)
            {
                agents_ = new Agent[RVOSimulator.Instance.agents_.Count];

                for (int i = 0; i < agents_.Length; ++i)
                {
                    agents_[i] = RVOSimulator.Instance.agents_[i];
                }

                agentTree_ = new AgentTreeNode[2 * agents_.Length];

                for (int i = 0; i < agentTree_.Length; ++i)
                {
                    agentTree_[i] = new AgentTreeNode();
                }
            }

            if (agents_.Length != 0)
            {
                BuildAgentTreeRecursive(0, agents_.Length, 0);
            }



        }

        /// <summary>
        /// 构建代理树递归
        /// </summary>
        /// <param name="begin"></param>
        /// <param name="end"></param>
        /// <param name="node"></param>
        internal void BuildAgentTreeRecursive(int begin, int end, int node)
        {
            agentTree_[node].begin = begin;
            agentTree_[node].end = end;
            agentTree_[node].minCoord = agents_[begin].position_;
            agentTree_[node].maxCoord = agents_[begin].position_;

            for (int i = begin + 1; i < end; i++)
            {
                agentTree_[node].maxCoord.x = Mathf.Max(agentTree_[node].maxCoord.x, agents_[i].position_.x);
                agentTree_[node].minCoord.x = Mathf.Min(agentTree_[node].minCoord.x, agents_[i].position_.x);
                agentTree_[node].maxCoord.y = Mathf.Max(agentTree_[node].maxCoord.y, agents_[i].position_.y);
                agentTree_[node].minCoord.y = Mathf.Min(agentTree_[node].minCoord.y, agents_[i].position_.y);
                agentTree_[node].maxCoord.z = Mathf.Max(agentTree_[node].maxCoord.z, agents_[i].position_.z);
                agentTree_[node].minCoord.z = Mathf.Min(agentTree_[node].minCoord.z, agents_[i].position_.z);
            }

            if (end - begin  > Math.RVO3D_MAX_LEAF_SIZE)
            {
                int coord;

                if (agentTree_[node].maxCoord.x - agentTree_[node].minCoord.x > agentTree_[node].maxCoord.y - agentTree_[node].minCoord.y 
                    && agentTree_[node].maxCoord.x - agentTree_[node].minCoord.x > agentTree_[node].maxCoord.z - agentTree_[node].minCoord.z)
                {
                    coord = 0;
                }
                else if (agentTree_[node].maxCoord.y - agentTree_[node].minCoord.y > agentTree_[node].maxCoord.z - agentTree_[node].minCoord.y)
                {
                    coord = 1;
                }
                else
                {
                    coord = 2;
                }

                float splitValue = 0.5f * (agentTree_[node].maxCoord[coord] + agentTree_[node].minCoord[coord]);

                int left = begin;
                int right = end;

                while (left < right)
                {
                    while (left < right && agents_[left].position_[coord] < splitValue)
                    {
                        ++left;
                    }

                    while (right > left && agents_[right - 1].position_[coord] >= splitValue)
                    {
                        --right;
                    }

                    if (left < right)
                    {
                        Agent t = agents_[left];
                        agents_[left] = agents_[right - 1];
                        agents_[right - 1] = t;

                        //Math.Swap(ref agents_[left], ref agents_[right - 1]);
                        ++left;
                        --right;
                    }
                }

                int leftSize = left - begin;

                if (leftSize ==0)
                {
                    ++leftSize;
                    ++left;
                    ++right;
                }

                agentTree_[node].left = node + 1;
                agentTree_[node].right = node + 2 * leftSize;

                BuildAgentTreeRecursive(begin, left, agentTree_[node].left);
                BuildAgentTreeRecursive(left, end, agentTree_[node].right);
            }
        }

        /// <summary>
        /// 计算指定Agent的Agent邻居
        /// </summary>
        /// <param name="agent"></param>
        /// <param name="rangeSq"></param>
        internal void ComputeAgentNeighbors(Agent agent, float rangeSq)
        {
            QueryAgentTreeRecursive(agent, rangeSq, 0);
        }

        /// <summary>
        /// 查询代理树递归
        /// </summary>
        /// <param name="agent"></param>
        /// <param name="rangeSq"></param>
        /// <param name="node"></param>
        internal void QueryAgentTreeRecursive(Agent agent, float rangeSq, int node)
        {
            if (agentTree_[node].end - agentTree_[node].begin <= Math.RVO3D_MAX_LEAF_SIZE)
            {
                for (int i = agentTree_[node].begin; i < agentTree_[node].end; i++)
                {
                    agent.InsertAgentNeighbor(agents_[i], rangeSq);
                }
            }
            else
            {
                float distSqLeft = Math.Square(Mathf.Max(0.0f, agentTree_[agentTree_[node].left].minCoord.x - agent.position_.x))
                    + Math.Square(Mathf.Max(0.0f, agent.position_.x - agentTree_[agentTree_[node].left].maxCoord.x))
                    + Math.Square(Mathf.Max(0.0f, agentTree_[agentTree_[node].left].minCoord.y - agent.position_.y))
                    + Math.Square(Mathf.Max(0.0f, agent.position_.y - agentTree_[agentTree_[node].left].maxCoord.y))
                    + Math.Square(Mathf.Max(0.0f, agentTree_[agentTree_[node].left].minCoord.z - agent.position_.z))
                    + Math.Square(Mathf.Max(0.0f, agent.position_.z - agentTree_[agentTree_[node].left].maxCoord.z));

                float distSqRight = Math.Square(Mathf.Max(0.0f, agentTree_[agentTree_[node].right].minCoord.x - agent.position_.x))
                    + Math.Square(Mathf.Max(0.0f, agent.position_.x - agentTree_[agentTree_[node].right].maxCoord.x))
                    + Math.Square(Mathf.Max(0.0f, agentTree_[agentTree_[node].right].minCoord.y - agent.position_.y))
                    + Math.Square(Mathf.Max(0.0f, agent.position_.y - agentTree_[agentTree_[node].right].maxCoord.y))
                    + Math.Square(Mathf.Max(0.0f, agentTree_[agentTree_[node].right].minCoord.z - agent.position_.z))
                    + Math.Square(Mathf.Max(0.0f, agent.position_.z - agentTree_[agentTree_[node].right].maxCoord.z));

                if (distSqLeft < distSqRight)
                {
                    if (distSqLeft < rangeSq)
                    {
                        QueryAgentTreeRecursive(agent, rangeSq, agentTree_[node].left);
                        if (distSqRight < rangeSq)
                        {
                            QueryAgentTreeRecursive(agent, rangeSq, agentTree_[node].right);
                        }
                    }
                }
                else
                {
                    if (distSqRight < rangeSq)
                    {
                        QueryAgentTreeRecursive(agent, rangeSq, agentTree_[node].right);
                        if (distSqLeft < rangeSq)
                        {
                            QueryAgentTreeRecursive(agent, rangeSq, agentTree_[node].left);
                        }
                    }
                }
            }
        }
    }
}

